import os
import numpy as np
from tensorflow.keras.applications import VGG16
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications.vgg16 import preprocess_input
from tensorflow.keras.models import Model
from sklearn.neighbors import NearestNeighbors
import pickle
from tqdm import tqdm
import cv2


def extract_features(img_path, model):
    """从单张图片提取特征"""
    img = image.load_img(img_path, target_size=(224, 224))
    x = image.img_to_array(img)
    x = np.expand_dims(x, axis=0)
    x = preprocess_input(x)
    features = model.predict(x)
    return features.flatten()


def train_model(dataset_path='static/dataset', model_save_path='models'):
    """训练模型并保存"""
    # 加载预训练的VGG16模型，去掉最后一层
    base_model = VGG16(weights='imagenet')
    model = Model(inputs=base_model.input, outputs=base_model.get_layer('fc1').output)

    # 确保模型保存目录存在
    os.makedirs(model_save_path, exist_ok=True)

    # 获取所有图片路径
    image_paths = []
    for root, _, files in os.walk(dataset_path):
        for file in files:
            if file.lower().endswith(('.png', '.jpg', '.jpeg')):
                image_paths.append(os.path.join(root, file))

    # 提取所有图片特征
    features = []
    for img_path in tqdm(image_paths, desc="Extracting features"):
        feat = extract_features(img_path, model)
        features.append(feat)

    features = np.array(features)

    # 训练最近邻模型
    neighbors = NearestNeighbors(n_neighbors=5, algorithm='brute', metric='cosine')
    neighbors.fit(features)

    # 保存模型和图片路径
    with open(os.path.join(model_save_path, 'neighbors.pkl'), 'wb') as f:
        pickle.dump(neighbors, f)

    with open(os.path.join(model_save_path, 'image_paths.pkl'), 'wb') as f:
        pickle.dump(image_paths, f)

    print(f"Model trained and saved to {model_save_path}")


if __name__ == '__main__':
    train_model()